#!/usr/bin/env python
# -*- coding: utf-8 -*-

import asyncio
import os
from typing import Dict, Any, List, Optional, Union

import litellm
from litellm import acompletion

from app.core.config import settings
from app.core.exceptions import LLMError
from app.core.logging import get_logger, LoggerMixin

# 创建日志记录器
logger = get_logger("llm")


class LLMManager(LoggerMixin):
    """LLM管理器
    
    使用LiteLLM与语言模型交互
    """
    
    def __init__(self):
        """初始化LLM管理器"""
        self.api_key = settings.LITELLM_API_KEY
        self.model = settings.LITELLM_MODEL
        self.temperature = settings.LITELLM_TEMPERATURE
        self.max_tokens = settings.LITELLM_MAX_TOKENS
        
        # 设置LiteLLM API密钥
        litellm.api_key = self.api_key
        
        # 如果使用的是 Ollama 模型，设置 Ollama 服务 URL
        if self.model.startswith("ollama/"):
            ollama_base_url = settings.OLLAMA_BASE_URL
            
            # 直接设置模型配置
            litellm.set_verbose = True  # 启用详细日志
            
            # 设置模型配置
            model_name = self.model.split("/")[1]  # 获取模型名称，例如 "qwen2.5:latest"
            litellm.model_list = [
                {
                    "model_name": f"ollama/{model_name}",
                    "litellm_params": {
                        "api_base": ollama_base_url
                    }
                }
            ]
            
            # 设置 Ollama API 基础 URL
            os.environ["OLLAMA_API_BASE"] = ollama_base_url
            
            self.logger.info(f"使用 Ollama 服务: {ollama_base_url}")
            self.logger.info(f"Ollama 模型配置: {litellm.model_list}")
        
        self.logger.info(f"已初始化LLM管理器，模型: {self.model}")
    
    async def generate_text(
        self,
        prompt: str,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        system_message: Optional[str] = None
    ) -> str:
        """生成文本
        
        Args:
            prompt: 提示词
            temperature: 温度参数（可选）
            max_tokens: 最大令牌数（可选）
            system_message: 系统消息（可选）
            
        Returns:
            str: 生成的文本
            
        Raises:
            LLMError: 当生成失败时
        """
        try:
            # 准备消息
            messages = []
            
            # 添加系统消息（如果提供）
            if system_message:
                messages.append({"role": "system", "content": system_message})
            
            # 添加用户消息
            messages.append({"role": "user", "content": prompt})
            
            # 准备参数
            params = {
                "model": self.model,
                "messages": messages,
                "temperature": temperature if temperature is not None else self.temperature,
                "max_tokens": max_tokens if max_tokens is not None else self.max_tokens
            }
            
            # 如果使用的是 Ollama 模型，添加 API 基础 URL
            if self.model.startswith("ollama/"):
                params["api_base"] = settings.OLLAMA_BASE_URL
            
            # 调用LiteLLM
            self.logger.info(f"调用 LiteLLM，参数: {params}")
            response = await acompletion(**params)
            
            # 提取生成的文本
            generated_text = response.choices[0].message.content
            
            self.logger.debug(f"文本生成成功，长度: {len(generated_text)}")
            return generated_text
            
        except Exception as e:
            error_msg = f"文本生成失败: {str(e)}"
            self.logger.error(error_msg)
            raise LLMError(error_msg)
    
    async def generate_sql(
        self,
        question: str,
        context: str,
        db_type: str = "mysql",
        system_prompt: Optional[str] = None
    ) -> str:
        """生成SQL查询
        
        Args:
            question: 用户问题
            context: 上下文信息（DDL、相似问题等）
            db_type: 数据库类型
            system_prompt: 自定义系统提示词
            
        Returns:
            str: 生成的SQL查询
            
        Raises:
            LLMError: 当生成失败时
        """
        # 构建系统消息
        default_system_message = (
            f"你是一个专业的SQL生成助手，精通{db_type}语法。"
            "你的任务是根据用户的问题和提供的上下文生成准确的SQL查询。"
            "只返回SQL查询，不要包含任何解释或其他文本。"
            "确保生成的SQL语法正确，可以直接在数据库中执行。"
            "非常重要：你必须严格遵循提供的数据库模式(DDL)和示例问题-SQL对。"
            "不要编造不存在的表或字段。如果上下文中没有足够的信息，请基于最相似的示例进行合理推断。"
            "生成的SQL必须与数据库模式兼容，并且风格应该与示例SQL保持一致。"
        )
        
        # 使用自定义系统提示词（如果提供）
        system_message = system_prompt if system_prompt else default_system_message
        
        # 构建提示词
        prompt = (
            f"请根据以下问题生成{db_type}查询：\n\n"
            f"问题：{question}\n\n"
            f"上下文信息：\n{context}\n\n"
            f"请生成一个能够回答上述问题的{db_type} SQL查询。只返回SQL查询，不要包含任何解释或其他文本。"
            f"你必须严格遵循提供的数据库模式和示例SQL。不要编造不存在的表或字段。"
        )
        
        # 生成SQL
        sql = await self.generate_text(
            prompt=prompt,
            system_message=system_message,
            temperature=0.0  # 使用较低的温度以获得更确定性的结果
        )
        
        # 清理SQL（移除可能的解释文本）
        sql_lines = sql.strip().split('\n')
        clean_sql_lines = []
        
        # 查找SQL代码块
        in_code_block = False
        for line in sql_lines:
            if line.startswith('```'):
                in_code_block = not in_code_block
                continue
            if in_code_block or not line.startswith('```'):
                clean_sql_lines.append(line)
        
        # 如果没有找到代码块，使用原始文本
        if not clean_sql_lines and sql_lines:
            clean_sql_lines = sql_lines
        
        clean_sql = '\n'.join(clean_sql_lines).strip()
        
        return clean_sql
    
    async def correct_sql(
        self,
        question: str,
        error_sql: str,
        error_message: str,
        context: str,
        db_type: str = "mysql",
        system_prompt: Optional[str] = None
    ) -> str:
        """纠正SQL查询
        
        Args:
            question: 原始问题
            error_sql: 错误的SQL查询
            error_message: 错误消息
            context: 上下文信息（DDL、相似问题等）
            db_type: 数据库类型
            system_prompt: 自定义系统提示词
            
        Returns:
            str: 纠正后的SQL查询
            
        Raises:
            LLMError: 当纠正失败时
        """
        # 构建系统消息
        default_system_message = (
            f"你是一个专业的SQL修复助手，精通{db_type}语法。"
            "你的任务是根据错误信息修复SQL查询。"
            "只返回修复后的SQL查询，不要包含任何解释或其他文本。"
            "确保修复后的SQL语法正确，可以直接在数据库中执行。"
        )
        
        # 使用自定义系统提示词（如果提供）
        system_message = system_prompt if system_prompt else default_system_message
        
        # 构建提示词
        prompt = (
            f"请修复以下{db_type} SQL查询中的错误：\n\n"
            f"原始问题：{question}\n\n"
            f"错误的SQL查询：\n{error_sql}\n\n"
            f"错误信息：\n{error_message}\n\n"
            f"上下文信息：\n{context}\n\n"
            f"请提供修复后的{db_type} SQL查询。只返回SQL查询，不要包含任何解释或其他文本。"
        )
        
        # 生成修复后的SQL
        corrected_sql = await self.generate_text(
            prompt=prompt,
            system_message=system_message,
            temperature=0.0  # 使用较低的温度以获得更确定性的结果
        )
        
        # 清理SQL（移除可能的解释文本）
        sql_lines = corrected_sql.strip().split('\n')
        clean_sql_lines = []
        
        # 查找SQL代码块
        in_code_block = False
        for line in sql_lines:
            if line.startswith('```'):
                in_code_block = not in_code_block
                continue
            if in_code_block or not line.startswith('```'):
                clean_sql_lines.append(line)
        
        # 如果没有找到代码块，使用原始文本
        if not clean_sql_lines and sql_lines:
            clean_sql_lines = sql_lines
        
        clean_sql = '\n'.join(clean_sql_lines).strip()
        
        return clean_sql


# 全局LLM管理器实例
llm_manager = LLMManager()


async def get_llm_manager() -> LLMManager:
    """获取LLM管理器
    
    Returns:
        LLMManager: LLM管理器实例
    """
    return llm_manager